"""
此脚本实现yolox的coco格式转为yolov5的yolo格式,要求的格式如下
coco格式:
    data/
        annotations/

        test/
        train/
        val/
        label.txt

yolo格式:
    data/
        images/
            train/
            test/
            val/
        labels/
            train/
            test/
            val/

调用示例:
    test = COCODataset(data_dir='/home/zsz/projects/YOLOX/datasets/fire_datasets/',
                        json_file='voc07_val.json',
                        name='val2017')
    test.convert_cocotoyolo()


data:2022.03.23
"""

from asyncore import write
import os
import cv2
import numpy as np
from pycocotools.coco import COCO


class COCODataset:
    """
    COCO dataset class.
    """

    def __init__(
            self,
            data_dir,
            json_file="voc07_test.json",
            name="train2017",
            img_size=(416, 416),
            preproc=None,
    ):
        """
        COCO dataset initialization. Annotation data are read into memory by COCO API.
        Args:
            data_dir (str): dataset root directory
            json_file (str): COCO json file name
            name (str): COCO data name (e.g. 'train2017' or 'val2017')
            img_size (int): target image size after pre-processing
            preproc: data augmentation strategy
        """
        self.data_dir = data_dir
        self.json_file = json_file

        self.coco = COCO(os.path.join(self.data_dir, "annotations", self.json_file))
        self.ids = self.coco.getImgIds()
        self.class_ids = sorted(self.coco.getCatIds())
        cats = self.coco.loadCats(self.coco.getCatIds())
        self._classes = tuple([c["name"] for c in cats])  # label信息
        self.name = name
        self.img_size = img_size
        self.preproc = preproc

    def __len__(self):
        return len(self.ids)

    def load_anno(self, index):
        id_ = self.ids[index]
        anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=False)
        annotations = self.coco.loadAnns(anno_ids)

        im_ann = self.coco.loadImgs(id_)[0]
        width = im_ann["width"]
        height = im_ann["height"]

        # load labels
        valid_objs = []
        for obj in annotations:
            x1 = np.max((0, obj["bbox"][0]))
            y1 = np.max((0, obj["bbox"][1]))
            x2 = np.min((width - 1, x1 + np.max((0, obj["bbox"][2] - 1))))
            y2 = np.min((height - 1, y1 + np.max((0, obj["bbox"][3] - 1))))
            if obj["area"] > 0 and x2 >= x1 and y2 >= y1:
                obj["clean_bbox"] = [x1, y1, x2, y2]
                valid_objs.append(obj)
        objs = valid_objs
        num_objs = len(objs)

        res = np.zeros((num_objs, 5))

        for ix, obj in enumerate(objs):
            cls = self.class_ids.index(obj["category_id"])
            res[ix, 0:4] = obj["clean_bbox"]
            res[ix, 4] = cls

        return res

    def pull_item(self, index):
        id_ = self.ids[index]
        im_ann = self.coco.loadImgs(id_)[0]
        width = im_ann["width"]
        height = im_ann["height"]
        filename = im_ann['file_name']

        # load image and preprocess
        img_file = os.path.join(
            self.data_dir, self.name, filename
        )
        img = cv2.imread(img_file)
        assert img is not None

        # load anno
        res = self.load_anno(index)
        img_info = {'size': (height, width), 'img_name': filename}

        boxes = []
        for tmp in res:
            x, y, w, h = self.convert_box(size=[width, height], box=tmp)
            boxes.append([int(tmp[-1]), x, y, w, h])

        return img, boxes, img_info, id_

    def make_dirs_yolo(self):
        """在原先的coco路径中创建yolo格式的文件夹
        """
        img_dir = os.path.join(self.data_dir, 'images')
        label_dir = os.path.join(self.data_dir, 'labels')

        img_train_dir = os.path.join(self.data_dir, 'images', 'train')
        img_test_dir = os.path.join(self.data_dir, 'images', 'test')
        img_val_dir = os.path.join(self.data_dir, 'images', 'val')

        label_train_dir = os.path.join(self.data_dir, 'labels', 'train')
        label_test_dir = os.path.join(self.data_dir, 'labels', 'test')
        label_val_dir = os.path.join(self.data_dir, 'labels', 'val')
        # img_train_dir,img_test_dir,img_val_dir,\
        tmps = [img_dir, label_dir,
                label_train_dir, label_test_dir, label_val_dir]
        for tmp in tmps:
            if not os.path.exists(tmp):
                os.makedirs(tmp)

    def convert_cocotoyolo(self):
        self.make_dirs_yolo()
        import shutil
        if self.name == 'train':
            img_org_dir = os.path.join(self.data_dir, 'train')
            img_dir = os.path.join(self.data_dir, 'images', 'train')
            label_dir = os.path.join(self.data_dir, 'labels', 'train')
        elif self.name == 'test':
            img_org_dir = os.path.join(self.data_dir, 'test')
            img_dir = os.path.join(self.data_dir, 'images', 'test')
            label_dir = os.path.join(self.data_dir, 'labels', 'test')
        elif self.name == 'val':
            img_org_dir = os.path.join(self.data_dir, 'val')
            img_dir = os.path.join(self.data_dir, 'images', 'val')
            label_dir = os.path.join(self.data_dir, 'labels', 'val')
        else:
            assert False
        if os.path.exists(img_dir):
            shutil.rmtree(img_dir)
        shutil.copytree(img_org_dir, img_dir)
        for idx in self.ids:
            img, boxes, img_info, id_ = self.pull_item(idx)
            # print(idx, img.shape, boxes, img_info, id_)
            assert img_info['img_name'][-3:] in ['png', 'jpg']
            save_txt_path = os.path.join(label_dir, img_info['img_name'][:-4] + '.txt')
            self.write_txt(boxes, save_txt_path)

    def write_txt(self, notes, save_txt_path):
        """写box信息到txt,专用.
        """
        with open(save_txt_path, 'w') as f:
            for note in notes:
                note = [str(x) for x in note]
                tmp = ' '.join(note) + '\n'
                f.write(tmp)

    def __getitem__(self, index):
        """
        One image / label pair for the given index is picked up and pre-processed.

        Args:
            index (int): data index

        Returns:
            img (numpy.ndarray): pre-processed image
            padded_labels (torch.Tensor): pre-processed label data.
                The shape is :math:`[max_labels, 5]`.
                each label consists of [class, xc, yc, w, h]:
                    class (float): class index.
                    xc, yc (float) : center of bbox whose values range from 0 to 1.
                    w, h (float) : size of bbox whose values range from 0 to 1.
            info_img : tuple of h, w, nh, nw, dx, dy.
                h, w (int): original shape of the image
                nh, nw (int): shape of the resized image without padding
                dx, dy (int): pad size
            img_id (int): same as the input index. Used for evaluation.
        """
        img, res, img_info, img_id = self.pull_item(index)

        if self.preproc is not None:
            img, target = self.preproc(img, res, self.input_dim)
        return img, target, img_info, img_id

    def convert_box(self, size, box):
        """转化xmin,ymin,xmax,ymax变成x,y,w,h. 这里的x,y是中心点.
            size: 图片的尺寸[w,h]
        """
        b1, b3, b2, b4 = box[:4]
        img_w, img_h = size
        # 标注越界修正
        if b2 > img_w:
            b2 = img_w
        if b4 > img_h:
            b4 = img_h
        b = (b1, b2, b3, b4)
        box = b

        dw = 1. / (size[0])
        dh = 1. / (size[1])
        x = (box[0] + box[1]) / 2.0 - 1
        y = (box[2] + box[3]) / 2.0 - 1
        w = box[1] - box[0]
        h = box[3] - box[2]
        x = x * dw
        w = w * dw
        y = y * dh
        h = h * dh
        return x, y, w, h


if __name__ == '__main__':
    train_set = COCODataset(data_dir='/home/sw/algo-env/dataset/fda84404-4c6f-4287-9f83-fbf397488343/fire/',
                            json_file='voc07_train.json',
                            name='train')
    train_set.convert_cocotoyolo()
    test_set = COCODataset(data_dir='/home/sw/algo-env/dataset/fda84404-4c6f-4287-9f83-fbf397488343/fire/',
                           json_file='voc07_test.json',
                           name='test')
    test_set.convert_cocotoyolo()
    val_set = COCODataset(data_dir='/home/sw/algo-env/dataset/fda84404-4c6f-4287-9f83-fbf397488343/fire/',
                          json_file='voc07_val.json',
                          name='val')
    val_set.convert_cocotoyolo()
